## 报损数据采集

import re
import pandas as pd
import json
import time
import datetime
import pytz
from playwright.sync_api import Playwright, sync_playwright, expect
import psycopg2

# 数据库连接配置
DB_CONFIG = {
    'dbname': 'heytea_db',
    'user': 'heytea',  # 确保使用具有足够权限的用户
    'password': '666666',
    'host': '127.0.0.1',
    'port': '5432'
}

def insert_data_to_db(data):
    try:
        # 连接到数据库
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()

        # 插入数据的SQL语句
        insert_query = """
        INSERT INTO material_loss_record (
            accounting_quantity, accounting_unit_id, accounting_unit_name, accounting_unit_spec,
            adjust_id, adjust_store, confirmed_quantity, is_confirmed, partner_id, product_code,
            product_id, product_name, quantity, unit_id, unit_name, unit_spec, adjust_date,
            updated_at, created_at, user_id, reason_type, convert_accounting_quantity,
            product_type, bom_type, source, amount, loss_date, loss_time, loss_moment,
            stock_qty, tax_rate, price, material_number
        ) VALUES (
            %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
            %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
        )
        """

        # 遍历数据并插入到数据库
        for row in data:
            cursor.execute(insert_query, (
                row.get('accounting_quantity'), row.get('accounting_unit_id'), row.get('accounting_unit_name'),
                row.get('accounting_unit_spec'), row.get('adjust_id'), row.get('adjust_store'),
                row.get('confirmed_quantity'), row.get('is_confirmed'), row.get('partner_id'),
                row.get('product_code'), row.get('product_id'), row.get('product_name'),
                row.get('quantity'), row.get('unit_id'), row.get('unit_name'), row.get('unit_spec'),
                row.get('adjust_date'), row.get('updated_at'), row.get('created_at'), row.get('user_id'),
                row.get('reason_type'), row.get('convert_accounting_quantity'), row.get('product_type'),
                row.get('bom_type'), row.get('source'), row.get('amount'), row.get('loss_date'),
                row.get('loss_time'), row.get('loss_moment'), row.get('stock_qty'), row.get('tax_rate'),
                row.get('price'), row.get('material_number')
            ))

        # 提交事务
        conn.commit()
        print("数据插入成功")

    except Exception as e:
        print(f"插入数据时发生错误: {e}")

    finally:
        # 关闭数据库连接
        if cursor:
            cursor.close()
        if conn:
            conn.close()

# 门店配置
stores = {
    "东里": {
        "username": "J21013414",
        "password": "J@21013414"
    },
    "莲下": {
        "username": "J21013524",
        "password": "J@21013524"
    }    
}

def run(playwright: Playwright, store_name: str, day_start: str, day_end: str, insert_to_db: bool) -> None:
    # 获取门店登录信息
    store = stores[store_name]
    username = store["username"]
    password = store["password"]

    browser = playwright.chromium.launch(headless=False)
    context = browser.new_context()
    page = context.new_page()
    page.set_default_timeout(30000)
    page.goto("https://partner-store.heytea.com/")

    try:        
        page.get_by_placeholder("请输入账号").fill(username)
        page.get_by_placeholder("请输入密码").fill(password)
        page.get_by_role("button", name="登录").click()
    except:
        print("logined")


    page.get_by_role("link", name="报废单").first.click()
    page.get_by_placeholder("Start Time").click()
    page.get_by_title(day_start).first.locator("div").click()
    page.get_by_title(day_end).first.locator("div").click()
    
    page.get_by_placeholder("Start Time").click()
    page.get_by_title(day_start).first.locator("div").click()
    page.get_by_title(day_end).first.locator("div").click()
    # page.get_by_role("button", name="查询").click()

    # 选中状态为已提交
    page.locator("#react-select-2--value div").filter(has_text="全部状态").click()
    page.get_by_role("option", name="已提交").click()

    with page.expect_response(re.compile(r".*/adjust.*")) as response_info:
        page.get_by_role("button", name="查询").click()
    row_response = response_info.value
    row_count = row_response.json()['payload']['total']
    data = []
    

    # row_count = page.locator('//*[@id="client-container"]//table/tbody//a').count()
    # print(["长度",row_count])
    for tr in range(row_count):
        with page.expect_response(re.compile(r".*/product.*")) as response_info:
            page.locator('//*[@id="client-container"]//div[@class="panel"]//table//tbody/tr[%s]/td[4]/a'%(tr+2)).dblclick()
        response = response_info.value
        rows = response.json()['payload']['rows']
        data.extend(rows)      
        page.get_by_role("link", name="报废单").first.click()

    # try:
    #     while page.get_by_role("button", name="right").count() > 0:
    #     # for i in range(0,3):
    #     # for i in range(0,3):
    #         with page.expect_response(re.compile(r".*/inventory/log\?.*")) as response_info:
    #             page.get_by_role("button", name="right").click()
    #         response = response_info.value
            
    #         # rows = json.loads(response.json()['rows']).rows
    #         rows = response.json()['payload']['rows']
    #         data.extend(rows)
    # except:
    #     print("timeouts!")
        # print(json.loads(response.json()))

    # page.get_by_role("button", name="right").click()
    # page.get_by_text("117").click()

    # ---------------------

    print(data)

    for row in data:
        utc_datetime = datetime.datetime.strptime(row['created_at'], "%Y-%m-%dT%H:%M:%SZ")
        # 设置本地时区
        local_timezone = pytz.timezone('Asia/Shanghai')  # 以上海时区为例
        local_datetime = utc_datetime.replace(tzinfo=pytz.utc).astimezone(local_timezone)

        row['created_at'] = local_datetime.strftime("%Y-%m-%d %H:%M:%S")
        row['loss_date'] = "{}-{}-{}".format(local_datetime.year,local_datetime.month,local_datetime.day)
        row['loss_time'] = local_datetime.strftime("%Y-%m-%d %H:%M:%S")
        row['loss_moment'] = local_datetime.strftime("%H:00:00")
    df = pd.DataFrame(data) 
    df.to_excel(f"{store_name}报损{day_start}-{day_end}.xlsx", index=False)

    # 在生成Excel文件后，根据开关决定是否插入数据到数据库
    if insert_to_db:
        insert_data_to_db(data)

    context.storage_state(path="cway")
    context.close()
    browser.close()

with sync_playwright() as playwright:
    # 获取用户输入的日期范围
    day_start = input("请输入查询起始日期（格式：-MM-DD）：")
    day_end = input("请输入查询结束日期（格式：-MM-DD）：")
    # 获取用户是否要插入数据库的选项
    insert_to_db = input("是否要将数据插入数据库？(y/n): ").lower() == 'y'
    
    # 显示所有可选门店
    print("可选门店：")
    for i, store_name in enumerate(stores.keys(), 1):
        print(f"{i}. {store_name}")
    
    # 获取用户选择的门店
    store_choice = input("请选择要操作的门店编号：")
    store_names = list(stores.keys())
    selected_store = store_names[int(store_choice) - 1]
    
    # 只运行用户选择的门店
    run(playwright, selected_store, day_start, day_end, insert_to_db)